"""
简单移动平均线交叉策略示例
"""
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from yhlz_backtest.strategy import Strategy
from yhlz_backtest.feed import PandasFeed
from yhlz_backtest.broker import Broker
from yhlz_backtest.plugins import ReturnsAnalyzer, TradeMetricsAnalyzer


class SimpleMovingAverageStrategy(Strategy):
    """
    简单移动平均线交叉策略
    
    当短期移动平均线上穿长期移动平均线时买入
    当短期移动平均线下穿长期移动平均线时卖出
    """
    
    def __init__(self, feed, broker, instrument, short_window=5, long_window=20):
        """
        初始化策略
        
        Parameters
        ----------
        feed : Feed
            数据源对象
        broker : Broker
            经纪人对象
        instrument : str
            交易品种代码
        short_window : int, optional
            短期窗口大小，默认为5
        long_window : int, optional
            长期窗口大小，默认为20
        """
        super().__init__(feed, broker)
        self._instrument = instrument
        self._short_window = short_window
        self._long_window = long_window
        self._prices = []
        self._position = 0
    
    def on_start(self):
        """
        策略开始时调用
        """
        super().on_start()
        self._logger.info(f"开始运行移动平均线交叉策略: {self._instrument}, 短期={self._short_window}, 长期={self._long_window}")
    
    def on_bar_impl(self, bar_dict):
        """
        处理每个bar的数据
        
        Parameters
        ----------
        bar_dict : dict
            包含各个品种当前bar数据的字典
        """
        if self._instrument not in bar_dict:
            return
        
        bar = bar_dict[self._instrument]
        
        # 获取收盘价
        close = bar['close']
        self._prices.append(close)
        
        # 确保有足够的数据
        if len(self._prices) <= self._long_window:
            return
        
        # 计算短期和长期移动平均线
        short_ma = np.mean(self._prices[-self._short_window:])
        long_ma = np.mean(self._prices[-self._long_window:])
        
        # 获取前一个短期和长期移动平均线
        if len(self._prices) > self._long_window + 1:
            prev_short_ma = np.mean(self._prices[-self._short_window-1:-1])
            prev_long_ma = np.mean(self._prices[-self._long_window-1:-1])
            
            # 检查交叉
            if prev_short_ma <= prev_long_ma and short_ma > long_ma:
                # 短期上穿长期，买入信号
                if self._position <= 0:
                    # 如果当前没有持仓或者是空仓，买入
                    self._logger.info(f"买入信号: 短期MA({short_ma:.2f}) > 长期MA({long_ma:.2f})")
                    
                    # 如果有空仓，先平仓
                    if self._position < 0:
                        self.buy(self._instrument, abs(self._position))
                    
                    # 开多仓
                    self.buy(self._instrument, 1)
                    self._position = 1
            
            elif prev_short_ma >= prev_long_ma and short_ma < long_ma:
                # 短期下穿长期，卖出信号
                if self._position >= 0:
                    # 如果当前有多仓或者没有持仓，卖出
                    self._logger.info(f"卖出信号: 短期MA({short_ma:.2f}) < 长期MA({long_ma:.2f})")
                    
                    # 如果有多仓，先平仓
                    if self._position > 0:
                        self.sell(self._instrument, abs(self._position))
                    
                    # 开空仓
                    self.sell(self._instrument, 1)
                    self._position = -1
    
    def on_finish(self):
        """
        策略结束时调用
        """
        super().on_finish()
        
        # 平掉所有仓位
        if self._position > 0:
            self.sell(self._instrument, abs(self._position))
        elif self._position < 0:
            self.buy(self._instrument, abs(self._position))
        
        self._logger.info("策略运行完成，已平掉所有仓位")


def run_strategy(data_file, instrument, initial_cash=100000.0, commission=0.0):
    """
    运行策略
    
    Parameters
    ----------
    data_file : str
        数据文件路径
    instrument : str
        交易品种代码
    initial_cash : float, optional
        初始资金，默认为100000.0
    commission : float, optional
        佣金比例，默认为0.0
        
    Returns
    -------
    tuple
        (策略对象, 收益率分析器, 交易指标分析器)
    """
    # 读取数据
    data = pd.read_csv(data_file, parse_dates=['datetime'], index_col='datetime')
    
    # 创建数据源
    feed = PandasFeed()
    feed.add_data_frame(instrument, data)
    
    # 创建经纪人
    broker = Broker(initial_cash=initial_cash, commission=commission)
    
    # 创建策略
    strategy = SimpleMovingAverageStrategy(feed, broker, instrument)
    
    # 运行策略
    results = strategy.run()
    
    # 分析结果
    returns_analyzer = ReturnsAnalyzer()
    returns_stats = returns_analyzer.analyze(results)
    
    metrics_analyzer = TradeMetricsAnalyzer()
    metrics_stats = metrics_analyzer.analyze(results, broker.filled_orders)
    
    print("--- 收益率统计 ---")
    for key, value in returns_stats.items():
        if isinstance(value, float):
            if key in ('total_return', 'annual_return', 'max_drawdown'):
                print(f"{key}: {value:.2%}")
            else:
                print(f"{key}: {value:.4f}")
        else:
            print(f"{key}: {value}")
    
    print("\n--- 交易指标统计 ---")
    if 'trade_count' in metrics_stats:
        print(f"总交易次数: {metrics_stats['trade_count']}")
    
    if 'daily_mean' in metrics_stats:
        print(f"平均每日收益率: {metrics_stats['daily_mean']:.2%}")
    
    # 绘制图表
    fig1 = returns_analyzer.plot()
    fig2 = metrics_analyzer.plot()
    
    plt.show()
    
    return strategy, returns_analyzer, metrics_analyzer


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description='运行简单移动平均线交叉策略')
    parser.add_argument('--data', type=str, required=True, help='数据文件路径')
    parser.add_argument('--instrument', type=str, required=True, help='交易品种代码')
    parser.add_argument('--cash', type=float, default=100000.0, help='初始资金')
    parser.add_argument('--commission', type=float, default=0.0, help='佣金比例')
    
    args = parser.parse_args()
    
    run_strategy(args.data, args.instrument, args.cash, args.commission) 